package avl.non_recursion_2;

/**
 * AVL Tree 类 <br>
 * <p>
 * 因为 Node 的插入，或删除，可能涉及根节点的变化，所以不能写在 Node 类中。
 *
 * @author xing
 * @date 2023-01-10 13:46:28
 */
public class AVLTree<V extends Comparable<V>> {
    private Node<V> root;

    public AVLTree() {
    }

    public AVLTree(Node<V> root) {
        this.root = root;
    }

    public void deleteElement(V element) {
        deleteNode(new Node<>(element));
    }

    public void deleteNode(Node<V> deleteNode) {
        this.root = deleteNode(this.root, deleteNode);
    }

    public Node<V> deleteNode(Node<V> root, Node<V> deleteNode) {
        if (root == null || deleteNode == null) {
            return root;
        }
        // 找到要被删除节点
        Node<V> node = findNodeOrGetParentByRoot(root, deleteNode);
        // 找到要被删除节点的父节点
        Node<V> parent = node.getParent();
        // 融合后的节点
        // 获取融合后的节点
        Node<V> mergeNode = mergeTwoTreeByParent(node);
        // 如果 parent 为 null，表明删除的是根节点，直接返回 mergeNode
        if (parent == null) {
            return mergeNode;
        }
        // mergeNode 可能为 null，需要根据 deleteNode 判断位置
        if (parent.getElement().compareTo(deleteNode.getElement()) < 0) {
            parent.setRight(mergeNode);
        } else {
            parent.setLeft(mergeNode);
        }
        // 修改子节点后重新计算高度
        parent.reComputerHeight();
        if (mergeNode != null) {
            mergeNode.setParent(parent);
        }
        // 从该节点开始，往父节点判断平衡，后返回
        return balanceTree(parent);
    }

    private Node<V> mergeTwoTreeByParent(Node<V> parent) {
        return mergeTwoTree(parent.getLeft(), parent.getRight());
    }

    private Node<V> mergeTwoTree(Node<V> left, Node<V> right) {
        if (left == null) {
            return right;
        }
        if (right == null) {
            return left;
        }
        // 都不为 Null
        // 删除父节点前都是平衡的，且之后还要从下边移动的节点往上判断平衡，
        // 这里不考虑左右子树的高度，直接使用左子树最大值替代根节点
        Node<V> lRight = left;
        while (lRight.getRight() != null) {
            lRight = lRight.getRight();
        }
        // 左节点没有右子节点, 左节点做根节点
        if (lRight == left) {
            // 设置新根节点的 parent 为 null
            // 别忘了两个方向都要设置，先对 parent 设置
            // （原来的 parent (left.parent) 相当于被删除了，
            //   不必做 left.getParent.setLeft(null) 或 setRight(null)，
            //   就是如果不设 Null， 不知道 JVM 底层当垃圾回收）
            // left.getParent().setLeft(null);
            // left.getParent().setRight(null);
            // 因为下边 balanceTheNodeWithParent() 会用到 parent, 虽然这里不设置，下边用了也没问题，但没必要，稳一点
            left.setParent(null);
            // 融合两子树
            left.setRightAndAsParent(right);
            // 重新计算高度
            left.reComputerHeight();
            // 进对该点做平衡后返回
            return balanceTheNodeWithParent(left);
        } else {
            // 左节点有右子节点，取得左子树最大节点的 parent
            Node<V> lrParent = lRight.getParent();
            // 设置新根节点的 parent 为 null
            // 别忘了两个方向都要设置，这里 parent.right 的设置在下边
            lRight.setParent(null);
            // 别忘了添加上 lRight 可能存在的左子节点
            lrParent.setRight(lRight.getLeft());
            if (lRight.getLeft() != null) {
                lRight.getLeft().setParent(lrParent);
            }
            // 融合两子树
            lRight.setLeftAndAsParent(left);
            lRight.setRightAndAsParent(right);
            // 从 parent 往 lright 链式计算高度
            lrParent.chainedReComputerHeight();
            // 从被删节点(做为根节点的 lRight)的父节点开始，往上平衡到现在子树的根节点
            // 后返回
            return balanceTree(lrParent);
        }
    }

    public void insertElement(V element) {
        if (element == null) {
            return;
        }
        // 插入节点
        insertNode(new Node<V>(element));
    }

    public void insertNode(Node<V> newNode) {
        this.root = insertNode(this.root, newNode);
        // 根据新插入节点，平衡插入后的树
        this.root = balanceTree(newNode);
    }

    /**
     * 插入新节点 (不处理平衡) <br>
     * TODO 这里如果前面构造函数只有需要参数的，root 不可能为 null，
     *      本身只做插入，就不需要返回值，并在外边重新赋值，且不必判断 root 是否为 null
     *      root 为 null 的判断不写在本方法里，也可以不要返回值，但如果要复用本方法，可能出现误用，导致报错
     *
     * @param root    要被插入的树的根节点
     * @param newNode 要插入的节点
     * @return 插入后的新树的根节点
     */
    private Node<V> insertNode(Node<V> root, Node<V> newNode) {
        if (newNode == null) {
            return root;
        }
        if (root == null) {
            return newNode;
        }
        // 获取要插入节点位置的父节点
        Node<V> parent = findNodeOrGetParentByRoot(root, newNode);
        // 插入新节点
        parent.setLeftOrRightAndAsParent(newNode);
        // 链式修改高度
        newNode.chainedReComputerHeight();
        return root;
    }

    /**
     * 节点存在，返回，不存在，按插入算，返回要插入位置的 parent
     *
     * @param root
     * @param node
     * @return 如果 node (树中与 node 元素值相同的节点)存在，返回该节点，<br>.如果 node 不存在，视为要插入，返回要插入的位置的 parent
     */
    private Node<V> findNodeOrGetParentByRoot(Node<V> root, Node<V> node) {
        // 存放节点的父节点（不赋初始值下边的 return 在 IDEA 会报错）
        Node<V> parent = root;
        // 获取新节点的父节点
        while (root != null) {
            parent = root;
            if (root.getElement().compareTo(node.getElement()) < 0) {
                root = root.getRight();
            } else if (root.getElement().compareTo(node.getElement()) > 0) {
                root = root.getLeft();
            } else {
                // 找到了
                // 如果要找到的也返回 parent，就在这里 return root.getparent(),
                break;
            }
        }
        return parent;
    }

    /**
     * 通过引起树变化的节点，平衡树, （该节点不会是不平衡点）<br>
     * 从 node 节点，往父节点方向，判断平衡，以及做平衡
     *
     * @param node 树中的某个节点
     * @return 平衡后的树的根节点
     */
    private Node<V> balanceTree(Node<V> node) {
        // node 不会是不平衡点，但是它的 parent 可能为 null，
        // 会导致下边 node = parent 后 为 null，从而报错 (不能 node.getParent())
        // Node<V> parent = node.getParent();
        Node<V> parent = node;
        do {
            node = parent;
            node = balanceTheNodeWithParent(node);
        } while ((parent = node.getParent()) != null);
        return node;
    }

    /**
     * 仅对 当前节点，判断该节点是否为平衡点，不平衡则旋转后返回新的根节点 <br>
     * 如果有父节点，旋转后，在方法外是不知道有没有平衡过的，会增加额外的不必要的操作。所以在方法内处理。<br>
     * 因为 父节点可能为 null，所以依然要返回一个平衡后的根节点。以及说，让人家知道你平衡后根节点在哪（是哪个）。
     *
     * @param root
     * @return
     */
    private Node<V> balanceTheNodeWithParent(Node<V> root) {
        // 正常来说，平衡树，变得只可能是根节点 root，而不会是 parent，
        // 所以这里 parent 不用作为参数，也不必关心返回
        final Node<V> parent = root.getParent();
        // 出现不平衡的基本条件，子节点高度起码为 2，即 root 的高度大于 2
        if (root.getHeight() > 2) {
            // 左子树高2，不平衡（一定会存在第二节点、第三节点）
            if (getHeight(root.getLeft()) - getHeight(root.getRight()) == 2) {
                Node<V> left = root.getLeft();
                // 根据高度判断第三节点
                if (getHeight(left.getLeft()) < getHeight(left.getRight())) {
                    // LR
                    root = doubleWithLeft(root);
                } else {
                    // 让相等的情况，也走更简单的右旋
                    // LL
                    root = rightRotate(root);
                }
            }
            // 右子树高2，不平衡
            else if (getHeight(root.getRight()) - getHeight(root.getLeft()) == 2) {
                Node<V> right = root.getRight();
                // 根据高度判断第三节点
                if (getHeight(right.getRight()) < getHeight(right.getLeft())) {
                    // RL
                    root = doubleWithRight(root);
                } else {
                    // 让相等的情况，也走更简单的左旋
                    // RR
                    root = leftRotate(root);
                }
            } else {
                // 平衡，直接返回
                return root;
            }
            // 不平衡处理完毕，重设与原父节点的关系
            if (parent != null) {
                parent.setLeftOrRightAndAsParent(root);
                // 修改完毕，从父节点链式计算高度
                parent.chainedReComputerHeight();
            }
        }
        return root;
    }

    /**
     * 左子树双旋
     */
    private Node<V> doubleWithLeft(Node<V> root) {
        Node<V> left = root.getLeft();
        // 在下边 left 中会重新设置 root.left 且其中用不到 root.left.parent
        // 这里不需要设置作为 parent, (root.setLeftAndAsParent())
        root.setLeft(leftRotate(left));
        // 这里 root 会在 rightRotate 后 修改高度，这里不再计算修改
        // 其实这里下边还会对 现在 root.left 重新计算，
        // 前面 leftRotate 里计算过了，浪费了一次计算，不想浪费就拆开单独写这里的逻辑，
        return rightRotate(root);
    }

    /**
     * 右子树双旋
     *
     * @param root 不平衡点
     * @return 平衡后的根节点
     */
    private Node<V> doubleWithRight(Node<V> root) {
        Node<V> right = root.getRight();
        // 在下边 left 中会重新设置 root.right 且其中用不到 root.right.parent
        // 这里不需要设置修改作为 parent (root.setRightAndAsParent())
        root.setRight(rightRotate(right));
        // 这里 root 会在 leftRotate 后 修改高度，这里不再计算修改
        return leftRotate(root);
    }

    /**
     * 右旋
     *
     * @param root 不平衡点
     * @return 平衡后的根节点
     */
    private Node<V> rightRotate(Node<V> root) {
        // 这里不做 root 与 root.left 为 null 的判断，不然出错了都不知道哪错了，或者自行抛异常
        Node<V> left = root.getLeft();
        root.setLeftAndAsParent(left.getRight());
        root.reComputerHeight();

        //下边的代码也可以这么写：
        //作用：
        //1. 另一个旋转也这样写，那么 doubleWithLeft 或 doubleWithRight 中的 set 可以省略，
        //   缺点：缺点是这里下边用的 setLeftOrRightAndAsParent()，也就是要判断，但前面二者不用判断，
        //        而且对于 double 总共要判断两次
        //2. balanceTheNodeWithParent() 里不再需要 parent 以及对 parent 的处理
        //   与第一条结合，这里原来 也是要对 parent 判断，做 setLeftOrRightAndAsParent()
        //   对于 leftRotate 与 rightRotate ，算是移过来了，但对于双旋，两种情况，是多了一次判断的，
        //   嗯，，唯一的好处就是不会出现 parent 忘了 链式计算的情况
        //   如果并不关心这点损失，确实可以这么做
        //Node<V> parent = root.getParent();
        //left.setParent(parent);
        //if (parent != null) {
        //    parent.setLeftOrRightAndAsParent(left);
        //}
        //left.setRightAndAsParent(root);
        //root.chainedReComputerHeight();

        left.setRightAndAsParent(root);
        left.reComputerHeight();
        // 别忘了 parent 的设置，由于是可能的根节点，所以设为 null，
        // 嗯，其实可以是 left.setParent(parent);
        // 这样，在 double 双旋方法中，中间就完全不需要使用 AsParent 的方法了，只 set 就非常够了（本来不改也够）
        // left 的 parent 也就是 root，在前面已经修改过 left 了，这里不必再设置，主需要关心它的 parent 的设置就可以了
        left.setParent(null);

        return left;
    }

    /**
     * 左旋
     *
     * @param root 不平衡点
     * @return 平衡后的根节点
     */
    private Node<V> leftRotate(Node<V> root) {
        // 这里不做 root 与 root.right 为 null 的判断，不然出错了都不知道哪错了，或者自行抛异常
        Node<V> right = root.getRight();
        root.setRightAndAsParent(right.getLeft());
        right.setLeftAndAsParent(root);
        // 注意顺序，现在 right 的 left 为 root，要自底向上修改计算高度
        // 或者前面 root 在 set 后就重新计算
        // TODO 可以将 重新计算写在 set 方法里，隐式自动修改自身高度属性。
        //      但是不写里边更灵活，有时候能省下一步重新计算，但也更容易出错，
        //      这里是 demo ，所以分开了，为了更好的表现出来，让功能浮于表面
        root.reComputerHeight();
        right.reComputerHeight();
        // 别忘了 parent 的设置，
        // right 的 parent 也就是 root，在前面已经修改过 right 了，这里不必再设置，主需要关心它的 parent 的设置就可以了
        right.setParent(null);
        return right;
    }

    /**
     * 获取 节点 高度，为 null 时返回 0
     *
     * @return 节点高度，如果为 null 返回 0
     */
    private int getHeight(Node<V> node) {
        return node == null ? 0 : node.getHeight();
    }

    public Node<V> getRoot() {
        return this.root;
    }
}
